import os
import shutil
from datetime import datetime
from PyQt5.QtCore import QThread, pyqtSignal, QSize, Qt
from PyQt5.QtWidgets import QWidget, QFileDialog, QListWidgetItem
from loguru import logger
from paddleocr import PaddleOCR
from qfluentwidgets import StateToolTip
from qfluentwidgets import FluentIcon as FIF

from config import cfg
from ui import ui_classification
from PyQt5.QtGui import QPixmap

from qt_tools.clip_predict import load_clip_model, clip_predict
from qt_tools.infobar import createWarningInfoBar, createSuccessInfoBar, createErrorInfoBar
from qt_tools.panddleocr import paddle_ocr_det_text


class Predict_Auto_Thread(QThread):
    response = pyqtSignal(dict)

    def __init__(self):
        super().__init__()
        self.pic_path_list = None
        self.base_path = None
        self.model_type = None
        self.clip_labels = None
        self.model = None

    def run(self) -> None:

        for index, i in enumerate(self.pic_path_list):
            try:
                if self.model_type == 'clip':
                    model, tokenizer, preprocess = self.model
                    result = clip_predict(os.path.join(self.base_path, i), self.clip_labels, model, tokenizer,
                                          preprocess)
                elif self.model_type == 'paddle':
                    result = paddle_ocr_det_text(self.model, os.path.join(self.base_path, i))

                if not result:
                    self.response.emit({'status': 'success', 'msg': '无识别结果', 'index': index})
                    continue

                self.response.emit({
                    'status': 'success',
                    'image': result,
                    'index': index})
            except Exception as e:
                self.response.emit({'status': 'errors', 'msg': str(e), 'index': index})

        self.response.emit({'status': 'all_success'})


class Predict_Thread(QThread):
    response = pyqtSignal(dict)

    def __init__(self):
        super().__init__()
        self.pic_path = None
        self.model_type = None
        self.clip_labels = None
        self.model = None

    def run(self) -> None:
        try:
            if self.model_type == 'clip':
                model, tokenizer, preprocess = self.model
                result = clip_predict(self.pic_path, self.clip_labels, model, tokenizer, preprocess)
            elif self.model_type == 'paddle':
                result = paddle_ocr_det_text(self.model, self.pic_path)

            self.response.emit({'status': 'success', 'image': result})
        except Exception as e:
            self.response.emit({'status': 'errors', 'msg': str(e)})


class Load_Model(QThread):
    response = pyqtSignal(dict)

    def __init__(self):
        super().__init__()
        self.model_type = None

    def run(self):
        try:
            path = cfg.get(cfg.model_downloadFolder)
            if not path:
                path = os.path.join(os.path.dirname(__file__), 'model')

            if self.model_type == 'clip':
                det = load_clip_model(path)
            elif self.model_type == 'paddle':
                det = PaddleOCR(lang="ch",
                                use_angle_cls=False,
                                det_model_dir=f"{path}\paddle\det\ch\ch_PP-OCRv4_det_infer",
                                cls_model_dir=f"{path}\paddle\cls\ch\ch_PP-OCRv4_cls_infer",
                                rec_model_dir=f"{path}\paddle\\rec\ch\ch_PP-OCRv4_rec_infer"
                                )

            self.response.emit({'status': 'success', 'model': det})
        except Exception as e:
            self.response.emit({'status': 'errors', 'msg': str(e)})


class Target_Class(QWidget):

    def __init__(self, ui_class=ui_classification.Ui_Form, parent=None):
        super().__init__(parent=parent)

        self.ui = ui_class()
        self.ui.setupUi(self)
        self.setObjectName('classification')

        self.arg_init()
        self.ui_init()
        self.btn_init()

    def arg_init(self):
        self.image_files_list = None  # 图片数据集列表
        self.datasets_base_path = None  # 存放图片的路径
        self.current_pic_index = None  # 当前选中的图片的索引
        self.label_path = None  # clip标签路径
        self.start_sign = False  # 开始识别标志
        self.use_model = None  # 使用中的模型
        self.model_list = []  # 模型的列表
        self.main_save_path = None  # 保存的路径

    def ui_init(self):
        self.ui.image_path_list.keyPressEvent = self.keyPressEvent

    def btn_init(self):
        self.ui.impor_dataset_btn.clicked.connect(self.import_datasets)
        self.ui.image_path_list.currentRowChanged.connect(self.show_images)

        self.ui.load_model_btn.clicked.connect(self.load_model)
        self.ui.load_clip_labels_btn.clicked.connect(self.import_clip_labels)
        self.ui.start_predict_btn.clicked.connect(self.start_task)

        self.ui.auto_tbtn.clicked.connect(self.auto_tbtn_event)
        self.ui.half_auto_tbtn.clicked.connect(self.half_auto_tbtn_event)
        self.ui.save_btn.clicked.connect(lambda: self.save_text_img(True))
        self.ui.skip_btn.clicked.connect(self.next_pic)

    def auto_tbtn_event(self):
        self.ui.half_auto_tbtn.setChecked(False)
        self.ui.start_predict_btn.setChecked(True)
        self.ui.start_predict_btn.setText('开启识别')

    def half_auto_tbtn_event(self):
        self.ui.auto_tbtn.setChecked(False)
        self.ui.start_predict_btn.setChecked(False)
        self.ui.start_predict_btn.setText('开启识别')

    def load_model(self):

        def response(infos):
            if infos['status'] == 'errors':
                self.stateTooltip.setContent('模型加载失败 ┭┮﹏┭┮')
                createErrorInfoBar(self, infos['msg'])
            else:
                self.stateTooltip.setContent('模型加载成功 😆')
                self.use_model = infos['model']

            self.stateTooltip.setState(True)

        if self.ui.start_predict_btn.isChecked():
            createWarningInfoBar(self, '请先关闭识别')
            return

        if self.ui.panddle_radio_btn.isChecked() and 'paddle' not in self.model_list:
            self.model_list.clear()
            self.model_list.append('paddle')
        elif self.ui.panddle_radio_btn.isChecked() and 'paddle' in self.model_list:
            createWarningInfoBar(self, '目前已经是paddle模型了')
            return

        if self.ui.clip_radio_btn.isChecked() and 'clip' not in self.model_list:
            self.model_list.clear()
            self.model_list.append('clip')
        elif self.ui.clip_radio_btn.isChecked() and 'clip' in self.model_list:
            createWarningInfoBar(self, '目前已经是clip模型了')
            return

        self.stateTooltip = StateToolTip('正在初始化加载模型', '客官请耐心等待哦~~', self)
        self.stateTooltip.move(0, 30)
        self.stateTooltip.show()

        self.load_model_thread = Load_Model()
        self.load_model_thread.model_type = self.model_list[0]
        self.load_model_thread.response.connect(response)
        self.load_model_thread.start()

    def import_datasets(self):

        folder_path = QFileDialog.getExistingDirectory(self, 'Select Folder',
                                                       f'{os.getcwd()}\image')
        if not folder_path:
            createWarningInfoBar(self, '请选择文件保存路径')
            return

        self.datasets_base_path = folder_path
        self.image_files_list = [f for f in os.listdir(folder_path) if
                                 f.lower().endswith(('png', 'jpg', 'jpeg', 'webp', 'bmp'))]

        if not self.image_files_list:
            createWarningInfoBar(self, '文件夹没有图片数据哦')
            return

        self.update_list_view()
        createSuccessInfoBar(self, '已导入图片数据')

    def import_clip_labels(self):
        folder_path, _ = QFileDialog.getOpenFileName(self, 'Select clip_label', filter='Text Files (*.txt)')
        if not folder_path:
            createWarningInfoBar(self, '请选择label文件')
            return
        self.label_path = folder_path
        createSuccessInfoBar(self, '导入完成')

    def update_list_view(self):
        self.ui.image_path_list.clear()
        for stand in self.image_files_list:
            item = QListWidgetItem(stand)
            self.ui.image_path_list.addItem(item)

        pixmap = QPixmap(os.path.join(self.datasets_base_path, self.image_files_list[0]))
        scaled_pixmap = pixmap.scaled(595, 529, Qt.AspectRatioMode.KeepAspectRatio)
        self.ui.PixmapLabel.setPixmap(scaled_pixmap)
        self.ui.image_path_list.setCurrentRow(0)

        self.t = datetime.now().strftime('%Y%m%d%H%M%S')

    def show_images(self):
        self.current_pic_index = self.ui.image_path_list.currentIndex().row()
        img_path = os.path.join(self.datasets_base_path, self.image_files_list[self.current_pic_index])
        current_pixmap = QPixmap(img_path)
        self.scaled_pixmap = current_pixmap.scaled(595, 529, Qt.AspectRatioMode.KeepAspectRatio)

        self.ui.PixmapLabel.setPixmap(self.scaled_pixmap)

        if self.ui.half_auto_tbtn.isChecked() and self.start_sign:
            self.start_predict_task()

    def start_task(self):
        if self.current_pic_index == None:
            createWarningInfoBar(self, '请先导入数据')
            self.ui.start_predict_btn.setChecked(False)
            return

        if self.use_model == None:
            createWarningInfoBar(self, '请先加载模型')
            self.ui.start_predict_btn.setChecked(False)
            return

        if not self.label_path or not os.path.exists(self.label_path):
            createWarningInfoBar(self, '请先新建label.txt或加载你的label文件')
            self.ui.start_predict_btn.setChecked(False)
            return

        if self.ui.auto_tbtn.isChecked():
            self.start_auto_predict_task()
            self.ui.start_predict_btn.setChecked(True)
        elif self.ui.half_auto_tbtn.isChecked() and self.ui.start_predict_btn.isChecked():
            self.start_sign = True
            self.start_predict_task()
        elif not self.ui.auto_tbtn.isChecked() and not self.ui.half_auto_tbtn.isChecked():
            createWarningInfoBar(self, '请选择识别模式')
            self.ui.start_predict_btn.setChecked(False)
            return

        if self.ui.start_predict_btn.isChecked() and not self.ui.auto_tbtn.isChecked():
            self.ui.start_predict_btn.setText('关闭识别')
        else:
            self.ui.start_predict_btn.setText('开启识别')
            self.start_sign = False

    def start_predict_task(self):
        def response(infos):
            logger.info(infos)
            if infos['status'] == 'errors':
                createErrorInfoBar(self, '图片识别失败')
                return
            elif infos['status'] == 'all_success':
                createSuccessInfoBar(self, '图片识别成功')
                return
            elif not infos['image']:
                createSuccessInfoBar(self, '图片识别成功,结果为空')
                return
            else:
                for i in infos['image']:
                    text, confiddence = i
                    self.ui.Confidence_label.setText(f'置信度: {str(float(confiddence))}')
                    self.ui.text_label_edit.setText(text)

        with open(self.label_path, encoding='utf8') as f:
            clip_labels = [i.strip() for i in f.readlines()]

        pic_path = os.path.join(self.datasets_base_path, self.image_files_list[self.current_pic_index])
        self.predict_thread = Predict_Thread()
        self.predict_thread.pic_path = pic_path
        self.predict_thread.model_type = self.model_list[0]
        self.predict_thread.model = self.use_model
        self.predict_thread.clip_labels = clip_labels
        self.predict_thread.response.connect(response)
        self.predict_thread.start()

    def start_auto_predict_task(self):
        def response(infos):
            if infos['status'] == 'all_success':
                createSuccessInfoBar(self, '图片识别成功')
                return

            self.current_pic_index = infos['index']
            if infos['status'] == 'errors':
                createErrorInfoBar(self, '图片识别失败')

                self.ui.image_path_list.setCurrentIndex(
                    self.ui.image_path_list.model().index(self.current_pic_index, 0))
                return

            if 'image' not in infos:
                createSuccessInfoBar(self, '图片识别成功,结果为空')
                self.ui.image_path_list.setCurrentIndex(
                    self.ui.image_path_list.model().index(self.current_pic_index, 0))
                return
            else:
                self.ui.image_path_list.setCurrentIndex(
                    self.ui.image_path_list.model().index(self.current_pic_index, 0))

                for i in infos['image']:
                    text, confiddence = i
                    self.ui.Confidence_label.setText(f'置信度: {str(float(confiddence))}')
                    self.ui.text_label_edit.setText(text)
                    if confiddence > 0.8:
                        self.save_text_img(True)
                    else:
                        self.save_text_img(False)

        if self.main_save_path is None:
            folder_path = QFileDialog.getExistingDirectory(self, 'Select Folder')
            if not folder_path:
                createWarningInfoBar(self, '请选择图片数据保存路径')
                return
            else:
                self.main_save_path = folder_path

        with open(self.label_path, encoding='utf8') as f:
            clip_labels = [i.strip() for i in f.readlines()]

        self.predict_thread = Predict_Auto_Thread()
        self.predict_thread.pic_path_list = self.image_files_list
        self.predict_thread.base_path = self.datasets_base_path
        self.predict_thread.model_type = self.model_list[0]
        self.predict_thread.model = self.use_model
        self.predict_thread.clip_labels = clip_labels
        self.predict_thread.response.connect(response)
        self.predict_thread.start()

    def save_text_img(self, save_type):
        """
        把标签结合到图片命名中保存
        """
        if self.main_save_path is None:
            self.main_save_path = cfg.get(cfg.pic_downloadFolder)
            if not self.main_save_path:
                createWarningInfoBar(self, '请先设置图片保存目录')
                return

        if self.current_pic_index is not None:
            text = self.ui.text_label_edit.text()
            if not text:
                createSuccessInfoBar(self, "当前图片识别结果为空")
                return
            if not save_type:
                text = '未知类别'

            if not os.path.exists(self.main_save_path + f'/{text}'):
                os.mkdir(self.main_save_path + f'/{text}')

            img_path = self.image_files_list[self.current_pic_index]
            total_img_path = os.path.join(self.datasets_base_path, img_path)
            save_path = os.path.join(self.main_save_path + f'/{text}', text + '_' + img_path)
            shutil.copy(total_img_path, save_path)

            createSuccessInfoBar(self, f"图片分类已保存到 {save_path}")
            if self.current_pic_index == len(self.image_files_list) - 1:
                createSuccessInfoBar(self, '图片已经遍历一轮')
            else:
                self.next_pic()

    def next_pic(self):
        self.ui.image_path_list.setCurrentRow(self.ui.image_path_list.currentIndex().row() + 1)

    def keyPressEvent(self, event):
        """
        键盘事件
        :param event:
        :return:
        """
        if event.key() == Qt.Key.Key_S:
            self.save_text_img(True)
